from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
import uvicorn
import asyncio
import cv2
import base64
import numpy as np
from typing import List

app = FastAPI()


# 模拟机器人控制
class RobotController:
    def __init__(self, session_id: str):
        self.session_id = session_id
        self.position = {"x": 0, "y": 0}
        self.speed = 1

    def move(self, direction: str):
        if direction == "forward":
            self.position["y"] += self.speed
        elif direction == "backward":
            self.position["y"] -= self.speed
        elif direction == "left":
            self.position["x"] -= self.speed
        elif direction == "right":
            self.position["x"] += self.speed
        return self.position


robots = {}


# WebSocket管理器
class ConnectionManager:
    def __init__(self):
        self.active_connections: List[WebSocket] = []
        self.robots: Dict[str, RobotController] = {}

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)
        session_id = str(id(websocket))
        self.robots[session_id] = RobotController(session_id)

    def disconnect(self, websocket: WebSocket):
        session_id = str(id(websocket))
        if session_id in self.robots:
            del self.robots[session_id]
        self.active_connections.remove(websocket)

    async def broadcast(self, message: str):
        for connection in self.active_connections:
            await connection.send_text(message)


manager = ConnectionManager()


# 视频捕获模拟
class VideoCapture:
    def __init__(self):
        self.cap = cv2.VideoCapture(0)  # 使用0表示默认摄像头
        if not self.cap.isOpened():
            # 如果摄像头不可用，使用模拟视频
            self.cap = None
            self.width, self.height = 640, 480
            self.count = 0

    async def get_frame(self):
        if self.cap and self.cap.isOpened():
            ret, frame = self.cap.read()
            if not ret:
                return None
        else:
            # 生成模拟视频帧
            frame = np.zeros((self.height, self.width, 3), np.uint8)
            cv2.putText(frame, f"Robot Camera {self.count}", (10, 30),
                        cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
            self.count += 1

        # 转换为JPEG并base64编码
        _, buffer = cv2.imencode('.jpg', frame)
        frame_base64 = base64.b64encode(buffer).decode('utf-8')
        return frame_base64


video_capture = VideoCapture()


# API路由
@app.get("/api/status")
async def get_status(session_id: str):
    if session_id in manager.robots:
        return {"position": manager.robots[session_id].position}
    return {"error": "Session not found"}


class MoveCommand(BaseModel):
    direction: str


@app.post("/api/move")
async def move_robot(command: MoveCommand, session_id: str):
    if session_id in manager.robots:
        position = manager.robots[session_id].move(command.direction)
        return {"message": f"Moved {command.direction}", "position": position}
    return {"error": "Session not found"}


class VoiceCommand(BaseModel):
    text: str


@app.post("/api/voice")
async def process_voice(command: VoiceCommand, session_id: str):
    # 简单的语音命令处理
    text = command.text.lower()
    direction = None

    if "forward" in text or "前进" in text:
        direction = "forward"
    elif "backward" in text or "后退" in text:
        direction = "backward"
    elif "left" in text or "左" in text:
        direction = "left"
    elif "right" in text or "右" in text:
        direction = "right"

    if direction and session_id in manager.robots:
        position = manager.robots[session_id].move(direction)
        return {"message": f"Executed voice command: {direction}", "position": position}
    return {"message": "Voice command not recognized"}


# WebSocket视频流
@app.websocket("/ws/video")
async def video_stream(websocket: WebSocket):
    await manager.connect(websocket)
    try:
        while True:
            frame = await video_capture.get_frame()
            if frame:
                await websocket.send_text(frame)
            await asyncio.sleep(0.05)  # 约20fps
    except WebSocketDisconnect:
        manager.disconnect(websocket)


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)